# -*- coding: utf-8 -*-
"""
轮动策略框架 | 邢不行 | 2024分享会
author: 邢不行
微信: xbx6660
"""
import pandas as pd

from Functions import *
from Evaluate import *
from Config import *
from joblib import Parallel, delayed
import sys
import warnings

warnings.filterwarnings('ignore')
pd.set_option('display.max_rows', 1000)
pd.set_option('expand_frame_repr', False)  # 当列太多时不换行

if __name__ == '__main__':
    # 动态读取config里面配置的shift_name脚本
    Shift = __import__('shift.%s' % shift_name, fromlist=('',))

    # =====取出Shift中配置的变量
    # 从Shift中取出 因子列表factor_list
    factor_list = Shift.factor_list
    # 从Shift中取出 覆盖参数列表params_list
    params_list = Shift.params_list
    # 从Shift中取出 回测持仓周期hold_period
    hold_period = Shift.hold_period
    # 从Shift中取出 offset
    offset = Shift.offset
    # 从Shift中取出 选策略数量select_num
    select_num = Shift.select_num
    # 从Shift中取出 资金曲线文件all_equity_file
    all_equity_file = Shift.all_equity_file
    # 从Shift中取出 子策略池base_strategy_dict
    base_strategy_dict = Shift.base_strategy_dict
    # 从Shift中取出 策略名stg_name
    stg_name = Shift.stg_name

    # =手动运行标识。手动运行的时候才会出画图，脚本调用直接跳过
    if_manual = True  # 默认手工运行
    if len(sys.argv) > 1:
        if_manual = False
        # 目前该脚本不支持遍历回测
        print('目前该脚本不支持遍历回测')
        exit()

    print('开始回测...')
    print('持仓周期：', hold_period)
    print('offset：', offset)
    print('因子列表：', factor_list)
    print('选策略数量：', select_num)

    # 将配置好的因子列表解析为 因子列表和参数列表
    factor_names = []
    factor_params = []
    factor_cols = []
    all_factor_list = []
    for _params in params_list:
        _f_list = []
        for i,_para in enumerate(_params):
            factor_name = factor_list[i][0]
            if_reverse = factor_list[i][1]
            weight = factor_list[i][3]
            _f_list.append((factor_name, if_reverse, _para, weight))
            factor_names.append(factor_name)
            factor_params.append(_para)
            factor_cols.append(factor_name + '_' + str(_para))
        all_factor_list.append(_f_list)
    factor_names = list(set(factor_names))
    factor_params = list(set(factor_params))
    factor_cols = list(set(factor_cols))

    # =====读取数据
    # 并行或串行，依次取出每个资金曲线的数据，最终生成一个总df，每个资金曲线当作币种进行处理
    multiply_process = False  # 是否并行。
    if multiply_process:
        equity_list = Parallel(n_jobs=n_jobs)(delayed(read_equity)(path, base_strategy_dict, hold_period, offset, benchmark) for path in all_equity_file)
    else:
        equity_list = []
        for path in all_equity_file:
            _equity = read_equity(path, base_strategy_dict, hold_period, offset, benchmark)
            equity_list.append(_equity)

    # 合并并整理数据
    base_equity = pd.concat(equity_list, axis=0)
    base_equity = base_equity.drop_duplicates(subset=['candle_begin_time', 'symbol'], keep='last')
    base_equity = base_equity.sort_values(['candle_begin_time', 'symbol']).reset_index(drop=True)  # 排序并重置索引

    # =====计算因子
    # 并行或串行，依次取出每个资金曲线的数据，计算因子
    multiply_process = False  # 是否并行。在测试的时候可以改成False，实际跑的时候改成True
    if multiply_process:
        df_list = Parallel(n_jobs=n_jobs)(delayed(calc_factors)(base_equity, symbol, factor_names, factor_params) for symbol in base_equity['symbol'].unique())
    else:
        df_list = []
        for symbol in base_equity['symbol'].unique():
            _equity = calc_factors(base_equity, symbol, factor_names, factor_params)
            df_list.append(_equity)
    # 合并并整理数据
    df = pd.concat(df_list, axis=0)
    if df[factor_cols].isnull().sum().max() == len(df):
        nan_df = df[factor_cols].isnull().sum()
        nan_df = nan_df[nan_df == nan_df.max()]
        print(f'{list(nan_df.index)}因子全部为空，请检查因子')
        exit()
    df = df.dropna(subset=factor_cols, axis=0)  # 删除因子值为空的数据
    df = df[(df['candle_begin_time'] >= start_date) & (df['candle_begin_time'] <= end_date)]  # 根据回测开始时间、结束时间截取数据
    if df.empty:
        print('日期配置有误，导致资金曲线为空，请检查日期及资金曲线文件')
        exit()
    df = df.sort_values(['candle_begin_time', 'symbol']).reset_index(drop=True)  # 排序

    # 保存一下轮动策略的中间数据用于轮动模拟器
    save_path = os.path.join(back_test_path, f'{stg_name}_轮动策略中间数据_{hold_period}_{offset}.pkl')
    save_inter_data(df, factor_cols, save_path)

    # 遍历得到因子列表
    all_equity = pd.DataFrame()
    for _factor_list in all_factor_list:
        _param = ''
        for factor_info in _factor_list:
            _param += f'{factor_info[2]}_'
        _param = _param[:-1]

        # ===根据因子选策略
        _df = df.copy()
        _df = select_strategy(_df, _factor_list, select_num)

        # ===计算资金曲线
        _df = cal_rot_net_value(_df, hold_period, spot_buy_crate, spot_sell_crate, swap_buy_crate, swap_sell_crate)

        # # 判断是否展示统计信息
        # if if_display_stat:
        #     # ===根据计算出的资金曲线数据得到轮动的评价指标
        #     display_stat_data(df, info_path, stg_name, hold_period, offset)

        # ===计算资金曲线
        # =创造空的时间周期表，用于填充不选币的周期
        _benchmark = get_benchmark(_df['candle_begin_time'].min(), _df['candle_begin_time'].max(), freq='1H')
        equity = pd.merge(_benchmark, _df[['candle_begin_time', '策略名称', '选币', '策略数量', 'label']], on='candle_begin_time', how='left')
        equity[['策略名称', '选币', '策略数量', 'label']] = equity[['策略名称', '选币', '策略数量', 'label']].fillna(method='ffill')
        equity['轮动涨跌幅'] = revise_data_length(_df['每小时涨跌幅'].sum(), len(equity))
        equity['_'] = equity['轮动涨跌幅'].fillna(method='bfill')
        if equity['_'].isnull().sum() > 0:
            index = equity[equity['_'].isnull()].index[0]
            equity = equity[equity.index < index].reset_index(drop=True)
        equity.drop('_', axis=1, inplace=True)
        equity['轮动资金曲线'] = (equity['轮动涨跌幅'] + 1).cumprod()

        cols = ['策略名称', '选币', '策略数量', 'label', '轮动涨跌幅', '轮动资金曲线']
        rename_dict = {col: f'{col}_{_param}' for col in cols}
        equity.rename(columns=rename_dict, inplace=True)
        all_equity = equity.copy() if all_equity.empty else pd.merge(all_equity, equity, on='candle_begin_time', how='left')
    # 由各个参数的轮动涨跌幅平均得到参数覆盖的轮动涨跌幅，在各个参数之间进行reb
    zdf_cols = [col for col in all_equity.columns if '轮动涨跌幅' in col]
    equity_cols = [col for col in all_equity.columns if '轮动资金曲线' in col]
    all_equity['轮动涨跌幅_reb'] = all_equity[zdf_cols].mean(axis=1)
    all_equity['轮动资金曲线_reb'] = (all_equity['轮动涨跌幅_reb'] + 1).cumprod()
    all_equity['轮动资金曲线_noreb'] = all_equity[equity_cols].mean(axis=1)  # 实盘暂时不支持noreb模式
    all_equity['轮动涨跌幅_noreb'] = all_equity['轮动资金曲线_noreb'].pct_change()  # 实盘暂时不支持noreb模式
    all_equity['轮动涨跌幅_noreb'].fillna(all_equity['轮动资金曲线_noreb'] - 1, inplace=True)

    # 保存资金曲线文件
    save_path = os.path.join(back_test_path, f'{stg_name}_轮动资金曲线_{hold_period}_{offset}.csv')
    all_equity.to_csv(save_path, encoding='gbk', index=False)

    # ===计算策略评价指标
    for _factor_list in all_factor_list:
        _param = ''
        for factor_info in _factor_list:
            _param += f'{factor_info[2]}_'
        _param = _param[:-1]
        rtn, year_return, month_return = strategy_evaluate(all_equity, net_col=f'轮动资金曲线_{_param}', pct_col=f'轮动涨跌幅_{_param}')
        print('=' * 20, _factor_list, '=' * 20)
        print(rtn, '\n\n分年收益率：\n', year_return, '\n\n分月收益率：\n', month_return)
        all_equity.rename(columns={'max2here': f'max2here_{_param}', 'dd2here': f'dd2here_{_param}'}, inplace=True)

    rtn_reb, year_return_reb, month_return_reb = strategy_evaluate(all_equity, net_col='轮动资金曲线_reb', pct_col='轮动涨跌幅_reb')
    print('=' * 20, 'reb', '=' * 20)
    print(rtn_reb, '\n\n分年收益率：\n', year_return_reb, '\n\n分月收益率：\n', month_return_reb)
    all_equity.rename(columns={'max2here': 'max2here_reb', 'dd2here': 'dd2here_reb'}, inplace=True)
    rtn_noreb, year_return_noreb, month_return_noreb = strategy_evaluate(all_equity, net_col='轮动资金曲线_noreb', pct_col='轮动涨跌幅_noreb')
    print('=' * 20, 'noreb', '=' * 20)
    print(rtn_noreb, '\n\n分年收益率：\n', year_return_noreb, '\n\n分月收益率：\n', month_return_noreb)
    all_equity.rename(columns={'max2here': 'max2here_noreb', 'dd2here': 'dd2here_noreb'}, inplace=True)

    # ===画图：
    if if_manual:
        # 画资金曲线图
        plot_ratation_net_value3(all_equity, factor_list, params_list)
    else:
        # 目前该脚本不支持遍历回测
        print('目前该脚本不支持遍历回测')
        exit()
